from io import BytesIO
import numpy as np
import proto_out.ndarray_pb2 as ndarray_pb2   #上面ndarray.proto编译成python
from proto_out.ndarray_pb2 import NDArray


def ndarray_to_proto(nda: np.ndarray) -> NDArray:
    """
    numpy转proto
    """
    nda_bytes = BytesIO()
    np.save(nda_bytes, nda, allow_pickle=False)
    return NDArray(ndarray=nda_bytes.getvalue())


def proto_to_ndarray(nda_proto: NDArray) -> np.ndarray:
    nda_bytes = BytesIO(nda_proto.ndarray)
    return np.load(nda_bytes, allow_pickle=False)



A = np.array([
    1, 2, 3, 4, 4,
    2, 3, 4, 5, 3,
    4, 5, 6, 7, 2,
    5, 6, 7, 8, 9,
    6, 7, 8, 9, 0 ]).reshape(5,5)
serialized_A = ndarray_to_proto(A)

print("Serialized message:")
print(serialized_A)

deserialized_A = proto_to_ndarray(serialized_A)
assert np.array_equal(A, deserialized_A)

